package com.easyto;

import java.lang.instrument.Instrumentation;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import javassist.*;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.instrument.Instrumentation;
import java.security.ProtectionDomain;

public class MyPreMainAgent {
    final static String source = "{\n" +
            "long begin = System.currentTimeMillis();\n" +
            "        Object result;\n" +
            "        try {\n" +
            "            result = ($w)%s$agent($$);\n" +
            "        }finally {\n" +
            "            long end = System.currentTimeMillis();\n" +
            "            System.out.println(\"agent cost : \" + (end-begin));\n" +
            "        }\n" +
            "        return ($r) result;\n" +
            "}\n";


    final static String voidSource = "{\n" +
            "long begin = System.currentTimeMillis();\n" +
            "        try {\n" +
            "            %s$agent($$);\n" +
            "        }finally {\n" +
            "            long end = System.currentTimeMillis();\n" +
            "            System.out.println(\"agent cost : \" + (end-begin));\n" +
            "        }\n" +
            "}\n";

    /**
     * java-agent 静态载入的情况 监控jvm运行情况
     * @param agentArgs
     * @param inst
     */
    //public static void premain(String agentArgs, Instrumentation inst) {
    //    System.out.println(">>>>>>> hello javaAgent");
    //
    //    Executors.newScheduledThreadPool(1).scheduleAtFixedRate(new Runnable() {
    //        @Override
    //        public void run() {
    //
    //            System.out.println(">>>>>>>>>> 开始监控 >>>>>>>>>>>> ");
    //            Metric.printMemoryInfo();
    //            Metric.printGCInfo();
    //        }
    //    }, 0, 3000, TimeUnit.MILLISECONDS);
    //}


    /**
     * 修改字节码，监控方法
     * @param agentArgs
     * @param instrumentation
     */
    public static void premain(String agentArgs, Instrumentation instrumentation) {
        System.out.println(">>>>>>> hello javaAgent");
        final String config = agentArgs;

        System.out.println(">>>>>>>>>>>>>>> config:   " + config);
        final ClassPool pool = new ClassPool();
        pool.appendSystemPath();

        //基于工具，在运行的时候修改class字节码，即动态插装
        instrumentation.addTransformer(new ClassFileTransformer() {
            @Override
            public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined, ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
                if (className == null || !className.replaceAll("/", ".").startsWith(config)) {
                    return null;
                }
                try {
                    className = className.replaceAll("/", ".");
                    CtClass ctClass = pool.getCtClass(className);
                    for (CtMethod ctMethod : ctClass.getDeclaredMethods()) {
                        newMethod(ctMethod);
                    }
                    return ctClass.toBytecode();
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return null;
            }
        });
    }

    private static CtMethod newMethod(CtMethod oldMethod) throws CannotCompileException, NotFoundException {
        CtMethod copy = CtNewMethod.copy(oldMethod, oldMethod.getDeclaringClass(), null);
        copy.setName(oldMethod.getName() + "$agent");
        oldMethod.getDeclaringClass().addMethod(copy);

        if (oldMethod.getReturnType().equals(CtClass.voidType)) {
            oldMethod.setBody(String.format(voidSource, oldMethod.getName()));
        } else {
            oldMethod.setBody(String.format(source, oldMethod.getName()));
        }
        return copy;
    }


}
